
import yfinance as yf
import pandas as pd
import numpy as np

# 读取历史K线数据
ticker = yf.Ticker('601919.SS')
df = ticker.history().drop(
    ["Volume", "Stock Splits"], axis=1)
print(f"股票数据为：{df}")

# 计算唐奇安上下轨线
n = 20
df['upperband'] = df['high'].rolling(n).max()
df['lowerband'] = df['low'].rolling(n).min()

# 判断是否突破唐奇安上下轨线
df.loc[df['close'] > df['upperband'], 'signal'] = 1
df.loc[df['close'] < df['lowerband'], 'signal'] = -1

# 进行交易
df['position'] = df['signal'].fillna(method='ffill') # 填充信号
df['position'].fillna(0, inplace=True) # 如果初始没有信号则默认空仓

df['returns'] = (df['close'] - df['close'].shift(1)) / df['close'].shift(1) # 计算每日收益率
df['strategy_returns'] = df['position'].shift(1) * df['returns'] # 计算每日策略收益率

df[['returns', 'strategy_returns']].dropna().cumsum().plot() # 绘制累计收益曲线

print(f"结果为：：：：：：：：：：：：")
# print(f"股票数据为：{df.head()}")